import sys
import typing

from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel, QPushButton, QFileDialog, QWidget
from PyQt5.QtGui import QPixmap, QPainter, QPen, QBrush
from PyQt5.QtCore import Qt, QRect
from loguru import logger
from Tools.DataStructs import RectInfo, PainterInfo


class YoloImageRectCache:
    def __init__(self, rect_list: typing.Tuple[RectInfo] = tuple()):
        self.rect_list_cache: typing.List[typing.Tuple[RectInfo]] = [rect_list, ]
        self.current_index = 0

    def append(self, rect_list: typing.Tuple[RectInfo]):
        # 说明要放弃之前current_index后面的了
        self.rect_list_cache = self.rect_list_cache[0:self.current_index + 1]
        self.rect_list_cache.append(rect_list)
        self.current_index += 1
        logger.info('current_index: %d, len:%d' % (self.current_index, len(self.rect_list_cache)))

    def undo(self):
        self.current_index = max(self.current_index - 1, 0)
        logger.info('current_index: %d, len:%d' % (self.current_index, len(self.rect_list_cache)))
        return self.get_current_rect_list()

    def redo(self):
        self.current_index = min(self.current_index + 1, len(self.rect_list_cache) - 1)
        logger.info('current_index: %d, len:%d' % (self.current_index, len(self.rect_list_cache)))
        return self.get_current_rect_list()

    def get_current_rect_list(self):
        if len(self.rect_list_cache) <= 0:
            return []
        if self.current_index < 0:
            raise IndexError('current_index[%d] < 0' % self.current_index)
        if self.current_index >= len(self.rect_list_cache):
            raise IndexError(
                'current_index[%d] >= len(rect_list_cache)[%d]' % (self.current_index, len(self.rect_list_cache)))
        temp_rect_list = list(self.rect_list_cache[self.current_index])
        logger.info('current_index: %d, len:%d, rect_list_len: %d' % (
            self.current_index, len(self.rect_list_cache), len(temp_rect_list)))
        return temp_rect_list

    def reset(self, rect_list: typing.Tuple[RectInfo] = tuple()):
        self.rect_list_cache = [rect_list, ]
        self.current_index = 0


class YoloImageWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.initUI()
        self.rect_list: typing.List[RectInfo] = []
        self.rect_list_cache: YoloImageRectCache = YoloImageRectCache()  # 配合撤销重做功能
        self.temp_rect = None
        self.painter_info = PainterInfo(Qt.red, '1')
        self.label_int = 1

    # public
    def set_painter_info(self, painter_info: PainterInfo, label_int):
        self.painter_info = painter_info
        self.label_int = label_int

    # private
    def initUI(self):
        self.setWindowTitle('YOLO Label Tool')
        self.setGeometry(100, 100, 800, 600)
        self.image_label = QLabel(self)
        self.image_label.resize(800, 600)
        self.image_label.move(0, 0)
        self.image_label.setStyleSheet("background-color: grey;")
        self.start_pos = None
        self.end_pos = None
        self.draw = False
        self.image = None
        self.pixmap = None
        self.temp_pixmap = None
        self.image_size = (800, 600)

    def openImage(self, fname, rect_list=None):
        if rect_list is None:
            rect_list = []
        self.image = QPixmap(fname)
        # image_sieze = self.image.size()
        self.pixmap = self.image.scaled(*self.image_size)
        self.image_label.setPixmap(self.pixmap)
        self.rest(rect_list)
        self.updateDrawing()

    def rest(self, rect_list=None):
        if rect_list is None:
            rect_list = []
        self.rect_list = rect_list
        self.rect_list_cache.reset(tuple(rect_list))

    def mousePressEvent(self, event):
        if self.image is None:
            return
        if event.button() == Qt.LeftButton:
            self.start_pos = event.pos()
            self.draw = True

    def mouseDoubleClickEvent(self, event):
        if event.button() != Qt.RightButton:
            return
        pose = event.pos()
        self.remove_rect_by_pose(pose)

    def remove_rect_by_pose(self, pose):
        for item in self.rect_list:
            if not item.rect.contains(pose):
                continue
            self.rect_list.remove(item)
            self.rect_list_cache.append(tuple(self.rect_list))
            self.updateDrawing()
            return

    def mouseMoveEvent(self, event):
        if self.image is None:
            return
        if event.buttons() == Qt.LeftButton and self.draw:
            self.end_pos = event.pos()
            self.temp_rect = QRect(self.start_pos, self.end_pos).normalized()
            self.updateDrawing()

    def mouseReleaseEvent(self, event):
        if self.image is None:
            return
        if event.button() == Qt.LeftButton:
            self.draw = False
            self.updateDrawing()
            self.rect_list.append(
                RectInfo(QRect(self.start_pos, self.end_pos).normalized(), self.label_int, self.painter_info))
            self.rect_list_cache.append(tuple(self.rect_list))
            self.temp_rect = None  # 鼠标临时绘制时候的矩形，绘制完成后重置为None
            self.start_pos = self.end_pos  # Update start position to end position

    def set_rect_list(self, rect_list):
        self.rect_list = rect_list
        self.rect_list_cache.append(tuple(self.rect_list))
        self.updateDrawing()

    def updateDrawing(self):
        if self.image is None:
            return
        self.temp_pixmap = self.pixmap.copy()
        painter = QPainter(self.temp_pixmap)
        brush = QBrush(Qt.white)  # 白色填充
        painter.setBrush(brush)
        painter.setOpacity(0.5)
        # draw saved info
        for rect_info in self.rect_list:
            painter.setPen(QPen(rect_info.painter_info.color, 2, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin))
            painter.drawRect(rect_info.rect)
            painter.drawText(rect_info.rect.center(), rect_info.painter_info.label)

        # draw temp rect
        if self.temp_rect is None:
            self.image_label.setPixmap(self.temp_pixmap)
            return
        painter.setPen(QPen(self.painter_info.color, 2, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin))
        painter.drawRect(self.temp_rect)
        painter.drawText(self.temp_rect.center(), self.painter_info.label)
        self.image_label.setPixmap(self.temp_pixmap)

    def saveAnnotation(self):
        # This method should be implemented to save the annotation in YOLO format
        pass

    def get_rect_info_json(self):
        return self.rect_list

    def get_yolo_format(self):
        return '\n'.join([rect_info.to_yolo_format(self.pixmap.size()) for rect_info in self.rect_list])

    @staticmethod
    def get_yolo_format_from_rect_list(yolo_image_widget, rect_list):
        return '\n'.join([rect_info.to_yolo_format(yolo_image_widget.pixmap.size()) for rect_info in rect_list])

    def undo(self):
        self.rect_list = self.rect_list_cache.undo()
        self.updateDrawing()

    def redo(self):
        self.rect_list = self.rect_list_cache.redo()
        self.updateDrawing()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = YoloImageWidget()
    window.openImage('/home/hello/图片/图片2.png')
    window.set_painter_info(PainterInfo(Qt.blue, '2'), 1)
    window.show()
    sys.exit(app.exec_())
